﻿namespace Util
{
    using System;
    using System.Linq;

    public static class KMeans
    {
        private static double[] WhenNumCentersGEValueLength(double[] value, int numCenters, out double cost)
        {
            cost = 0;
            double[] cen = new double[numCenters];
            for (int i = 0; i < value.Length; ++i)
            {
                cen[i] = value[i];
            }
            for (int i = value.Length; i < numCenters; ++i)
            {
                cen[i] = value.Last();
            }
            return cen;
        }

        /// <summary>
        /// Time complexity is O(value.Length * value.Length * numCenters) when numCenters >= 2.
        /// </summary>
        /// <param name="value">Used to find the K-Means centers, and may be changed (sorted) according to the algorithm needs.</param>
        /// <param name="numCenters">Number of K-means centers.</param>
        /// <param name="cost">The minimum cost given the optimum centers.</param>
        /// <returns>The centers, which is sorted.</returns>
        public static double[] KMeansOneDim(double[] value, int numCenters, out double cost)
        {
            Array.Sort(value);
            if (numCenters >= value.Length) return WhenNumCentersGEValueLength(value, numCenters, out cost);

            double[] distToFirst = new double[value.Length];
            distToFirst[0] = 0;
            for (int i = 1; i < value.Length; ++i)
            {
                distToFirst[i] = distToFirst[i - 1] + value[i] - value[0];
            }

            double[] distToLast = new double[value.Length];
            distToLast[value.Length - 1] = 0;
            for (int i = value.Length - 2; i >= 0; i--)
            {
                distToLast[i] = distToLast[i + 1] + value.Last() - value[i];
            }

            // DP table. f[i, j] represents the minimum cost when using i+1 centers where the last one is located at value[j].
            double[,] f = new double[numCenters, value.Length];

            // Record of the state transition path.
            int[,] fromIndex = new int[numCenters, value.Length];

            // Boundary condition: when using only 1 center and the center is placed at point value[j], the min cost
            for (int j = 0; j < value.Length; ++j)
            {
                f[0, j] = distToFirst.Last() - distToFirst[j] - (value.Length - j - 1) * (value[j] - value[0]) +  // The total distance from the right part to value[j]
                        distToLast.First() - distToLast[j] - j * (value.Last() - value[j]);  // The total distance from the left part to value[j]
            }

            // DP
            for (int i = 2; i <= numCenters; ++i)  // Using i centers
            {
                for (int j = i - 1; j < value.Length; ++j)  // The last center of these i centers is located at value[j]
                {
                    // The total distance from the right part to value[j]
                    double sj = distToFirst.Last() - distToFirst[j] - (value.Length - j - 1) * (value[j] - value[0]);

                    // Place the last two centers at the same place: value[j]
                    f[i - 1, j] = f[i - 2, j] - sj;
                    fromIndex[i - 1, j] = j;

                    int lowerBoundIndex = i - 2;
                    for (int k = i - 2; k < j; ++k)
                    {
                        if (k > i - 2 && value[k] == value[k - 1]) continue;
                        if (value[k] == value[j]) break;
                        // Now we must have value[k] < value[j]

                        // sk is the total distance from the right part to value[k]
                        double sk = distToFirst.Last() - distToFirst[k] - (value.Length - k - 1) * (value[k] - value[0]);

                        // bkj is the minimum total distance from value[k, k+1, ..., j] to value[k] or value[j]
                        double mid = value[k] + (value[j] - value[k]) / 2.0;
                        while (value[lowerBoundIndex] < mid) ++lowerBoundIndex;  // Amortized

                        // k, k+1, ..., lowerBoundIndex-1, lowerBoundIndex, mid+1, ..., j-1, j
                        // ------------------------------  -----------------------------------
                        //            left part                        right part
                        double distOfLeftPart = distToFirst[lowerBoundIndex - 1] - distToFirst[k] - (lowerBoundIndex - 1 - k) * (value[k] - value[0]);
                        double distOfRightPart = distToLast[lowerBoundIndex] - distToLast[j] - (j - lowerBoundIndex) * (value.Last() - value[j]);
                        double bkj = distOfLeftPart + distOfRightPart;

                        double cur = f[i - 2, k] - sk + bkj;
                        if (cur < f[i - 1, j])
                        {
                            f[i - 1, j] = cur;
                            fromIndex[i - 1, j] = k;
                        }
                    }
                    f[i - 1, j] += sj;
                }
            }

            cost = double.MaxValue;
            int resultIndex = -1;
            for (int j = numCenters - 1; j < value.Length; ++j)
            {
                if (f[numCenters - 1, j] < cost)
                {
                    cost = f[numCenters - 1, j];
                    resultIndex = j;
                }
            }

            double[] centers = new double[numCenters];
            centers[numCenters - 1] = value[resultIndex];
            for (int i = numCenters - 2; i >= 0; --i)
            {
                resultIndex = fromIndex[i + 1, resultIndex];
                centers[i] = value[resultIndex];
            }
            return centers;
        }

        /// <summary>
        /// Prerequisite: @value must be sorted, each element in @value must be different.
        /// @count[i] is the number of points whos value is @value[i].
        /// </summary>
        /// <param name="value"></param>
        /// <param name="count"></param>
        /// <param name="numCenters"></param>
        /// <param name="cost"></param>
        /// <returns></returns>
        public static double[] KMeansOneDimWithSortedCount(double[] value, int[] count, int numCenters, out double cost)
        {
            if (numCenters >= value.Length) return WhenNumCentersGEValueLength(value, numCenters, out cost);

            int[] sum = new int[count.Length];
            sum[0] = count[0];
            for (int i = 1; i < sum.Length; ++i)
            {
                sum[i] = sum[i - 1] + count[i];
            }

            double[] distToFirst = new double[value.Length];
            distToFirst[0] = 0;
            for (int i = 1; i < value.Length; ++i)
            {
                distToFirst[i] = distToFirst[i - 1] + (value[i] - value[0]) * count[i];
            }

            double[] distToLast = new double[value.Length];
            distToLast[value.Length - 1] = 0;
            for (int i = value.Length - 2; i >= 0; i--)
            {
                distToLast[i] = distToLast[i + 1] + (value.Last() - value[i]) * count[i];
            }

            // DP table
            double[,] f = new double[numCenters, value.Length];
            int[,] fromIndex = new int[numCenters, value.Length];

            // Boundary condition: when using only 1 center and the center is placed at point value[j], the min cost
            for (int j = 0; j < value.Length; ++j)
            {
                f[0, j] = distToFirst.Last() - distToFirst[j] - (sum.Last() - sum[j]) * (value[j] - value[0]) +  // The total distance from the right part to value[j]
                    distToLast.First() - distToLast[j] - (j == 0 ? 0 : sum[j - 1]) * (value.Last() - value[j]);  // The total distance from the left part to value[j]
            }

            // DP
            for (int i = 2; i <= numCenters; ++i)  // Using i centers
            {
                for (int j = i - 1; j < value.Length; ++j)  // The last center of these i centers is located at value[j]
                {
                    // The total distance from the right part to value[j]
                    double sj = distToFirst.Last() - distToFirst[j] - (sum.Last() - sum[j]) * (value[j] - value[0]);

                    // The left part of the total cost, initialized as we're placing the last two centers at the same place: value[j]
                    f[i - 1, j] = f[i - 2, j] - sj;
                    fromIndex[i - 1, j] = j;

                    int lowerBoundIndex = i - 2;
                    for (int k = i - 2; k < j; ++k)
                    {
                        // sk is the total distance from the right part to value[k]
                        double sk = distToFirst.Last() - distToFirst[k] - (sum.Last() - sum[k]) * (value[k] - value[0]);

                        // bkj is the minimum total distance from value[k, k+1, ..., j] to value[k] or value[j]
                        double mid = value[k] + (value[j] - value[k]) / 2.0;
                        while (value[lowerBoundIndex] < mid) ++lowerBoundIndex;  // Amortized O(1)

                        // k, k+1, ..., lowerBoundIndex-1, lowerBoundIndex, mid+1, ..., j-1, j
                        // ------------------------------  -----------------------------------
                        //            left part                        right part
                        double distOfLeftPart = distToFirst[lowerBoundIndex - 1] - distToFirst[k] - (sum[lowerBoundIndex - 1] - sum[k]) * (value[k] - value[0]);
                        double distOfRightPart = distToLast[lowerBoundIndex] - distToLast[j] - (sum[j - 1] - sum[lowerBoundIndex - 1]) * (value.Last() - value[j]);
                        double bkj = distOfLeftPart + distOfRightPart;

                        double cur = f[i - 2, k] - sk + bkj;
                        if (cur < f[i - 1, j])
                        {
                            f[i - 1, j] = cur;
                            fromIndex[i - 1, j] = k;
                        }
                    }
                    f[i - 1, j] += sj;
                }
            }

            cost = double.MaxValue;
            int resultIndex = -1;
            for (int j = numCenters - 1; j < value.Length; ++j)
            {
                if (f[numCenters - 1, j] < cost)
                {
                    cost = f[numCenters - 1, j];
                    resultIndex = j;
                }
            }

            double[] centers = new double[numCenters];
            centers[numCenters - 1] = value[resultIndex];
            for (int i = numCenters - 2; i >= 0; --i)
            {
                resultIndex = fromIndex[i + 1, resultIndex];
                centers[i] = value[resultIndex];
            }
            return centers;
        }
    }
}
